import wave
import pyaudio
import numpy as np
import matplotlib.pyplot as plt

#####################################################
# 读取.wav格式的音频文件，进行PSD功率谱密度分析
# 基于滑动时间窗口
#####################################################
#            参 数 
# 采样率：48KHz
# 时间窗口大小：1s (48000 points)
# step: 0.1s (4800 points)# #
#####################################################


#wavfile = "10KHz-stero.wav"
wavfile = "场景1.wav"
p = pyaudio.PyAudio()#创建PyAudio对象
wf=wave.open(wavfile,"rb")#打开wav
params = wf.getparams()#参数获取
nchannels, sampwidth, framerate, nframes = params[:4]

stream = p.open(format=p.get_format_from_width(sampwidth),
                channels=nchannels, 
                rate=framerate,
                output=True)#创建输出流
#读取完整的帧数据到str_data中，这是一个string类型的数据
str_data = wf.readframes(nframes)
wf.close()#关闭wave

#####2.将波形数据转换为数组
# N-1 一维数组，右声道接着左声道
wave_data = np.frombuffer(str_data, dtype=np.short)
#2-N N维数组
wave_data.shape = -1,2
#将数组转置为 N-2 目标数组
wave_data = wave_data.T

#########################################
wave_length = len(wave_data[0])
print("length of wave_data:%d" % wave_length) #440320

window_size=48000*1 # 采样率*滑动窗口大小
step = 4800
begin = 0
ending = begin+window_size

frequency_bin_min = 9900
frequency_bin_max = 10100
output = []
freqs = np.linspace(begin, ending, window_size)

while ending<wave_length:
    #print(len(wave_data[0][begin:ending]))
    _r = np.log10(np.abs(np.fft.rfft(wave_data[0][begin:ending]))**2/window_size)
    output.append(_r[frequency_bin_min:frequency_bin_max])
    begin += step
    ending += step

_out = np.array(output)
#print (_out.shape)
###################################################
#              draw heatmap figure                # 
###################################################

# plt.figure(figsize=(20,5))
# plt.matshow(_out, cmap = plt.cm.cool, vmin=_out.min(), vmax=_out.max())
# plt.colorbar()
# plt.show()
    
###################################################
#             draw 3D surface figure              # 
###################################################

fig = plt.figure()
ax3 = plt.axes(projection='3d')
xx = np.linspace(0,_out.shape[0],_out.shape[0])
yy = np.linspace(0,_out.shape[1],_out.shape[1])
X, Y = np.meshgrid(yy, xx)
ax3.plot_surface(X,Y,_out,cmap='rainbow')
plt.show()